use blake2::{Blake2b512, Digest};
use rand::RngCore;
use crate::error::{Error, Result};

/// Salamander obfuscator for QUIC packets
pub struct SalamanderObfuscator {
    key: Vec<u8>,
}

impl SalamanderObfuscator {
    pub fn new(password: impl AsRef<[u8]>) -> Self {
        Self {
            key: password.as_ref().to_vec(),
        }
    }

    pub fn obfuscate(&self, payload: &[u8]) -> Vec<u8> {
        let mut salt = [0u8; 8];
        rand::thread_rng().fill_bytes(&mut salt);

        let hash = self.compute_hash(&salt);

        let mut result = Vec::with_capacity(8 + payload.len());
        result.extend_from_slice(&salt);

        for (i, &byte) in payload.iter().enumerate() {
            result.push(byte ^ hash[i % 32]);
        }

        result
    }

    pub fn deobfuscate(&self, data: &[u8]) -> Result<Vec<u8>> {
        if data.len() < 8 {
            return Err(Error::InvalidPacket);
        }

        let salt = &data[0..8];
        let payload = &data[8..];

        let hash = self.compute_hash(salt);

        let mut result = Vec::with_capacity(payload.len());
        for (i, &byte) in payload.iter().enumerate() {
            result.push(byte ^ hash[i % 32]);
        }

        Ok(result)
    }

    fn compute_hash(&self, salt: &[u8]) -> [u8; 64] {
        let mut hasher = Blake2b512::new();
        hasher.update(&self.key);
        hasher.update(salt);
        let hash = hasher.finalize();

        let mut result = [0u8; 64];
        result.copy_from_slice(&hash);
        result
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_salamander_roundtrip() {
        let obfuscator = SalamanderObfuscator::new("test_password");
        let original = b"Hello, World! This is a test message.";

        let obfuscated = obfuscator.obfuscate(original);
        assert_ne!(&obfuscated[8..], original);
        assert_eq!(obfuscated.len(), original.len() + 8);

        let deobfuscated = obfuscator.deobfuscate(&obfuscated).unwrap();
        assert_eq!(&deobfuscated, original);
    }

    #[test]
    fn test_different_passwords() {
        let obf1 = SalamanderObfuscator::new("password1");
        let obf2 = SalamanderObfuscator::new("password2");
        let original = b"Test data";

        let obfuscated = obf1.obfuscate(original);
        
        // Should fail with wrong password
        let result = obf2.deobfuscate(&obfuscated);
        assert!(result.is_ok());
        assert_ne!(&result.unwrap(), original);
    }
}
